import torch
import sys
import numpy as np
import os
import random
import torch.nn.functional as F
from sklearn.metrics.pairwise import cosine_similarity
from torch.utils.data import DataLoader

sys.path.append(os.path.dirname(os.path.dirname(os.getcwd())))


def get_dataloader(data_folder, view_list = [1, 2, 3], num_class = 2, batch_size=32, num_workers=1, train_shuffle=True):
        
    num_view = len(view_list)
    
    labels_tr = np.loadtxt(os.path.join(data_folder, "labels_tr.csv"), delimiter=',')
    labels_te = np.loadtxt(os.path.join(data_folder, "labels_te.csv"), delimiter=',')
    labels_tr = labels_tr.astype(int)
    labels_te = labels_te.astype(int)
    
    data_tr_list = []
    data_te_list = []
    for i in view_list:
        data_tr_list.append(np.loadtxt(os.path.join(data_folder, str(i)+"_tr.csv"), delimiter=','))
        data_te_list.append(np.loadtxt(os.path.join(data_folder, str(i)+"_te.csv"), delimiter=','))
    
    num_tr = data_tr_list[0].shape[0]
    num_te = data_te_list[0].shape[0]
    
    data_mat_list = []
    for i in range(num_view):
        data_mat_list.append(np.concatenate((data_tr_list[i], data_te_list[i]), axis=0))
    
    labels = np.concatenate((labels_tr, labels_te)) 
    
    le = num_tr + num_te
    datasets = []
    for i in range(le):
        datasets.append([data_mat_list[0][i], data_mat_list[1][i], data_mat_list[2][i], labels[i]])
    
    random.seed(10)

    random.shuffle(datasets)
    
    tests = DataLoader(datasets[0:le//5], shuffle=False, 
                       num_workers=num_workers, batch_size=batch_size)  # 20% 测试集
    valids = DataLoader(datasets[le//5:le//4], shuffle=False, 
                        num_workers=num_workers, batch_size=batch_size)  # 5% 验证集
    trains = DataLoader(datasets[le//4:], shuffle=train_shuffle, 
                        num_workers=num_workers, batch_size=batch_size)  # 75% 训练集

    return trains, valids, tests